/*
 * NormMediaWriter.cpp
 *
 *  Created on: 2018年9月28日
 *      Author: zhengboyuan
 */

#include "NormMediaWriter.h"
#include "CodecIDConvert.h"
#include "MFormatUtil.h"
#include "TFileUtil.h"
#include "H26xNaluParser.h"
#include <assert.h>


namespace av
{

NormMediaWriter::NormMediaWriter(MediaFileFormat fileFormat, MediaWriter* writer):
	m_fileFormat(fileFormat),
	m_writer(writer),
	m_startVideoPts(-1),
	m_startAudioPts(-1)
{
}

NormMediaWriter::~NormMediaWriter()
{
	close();

	delete m_writer;
}

bool NormMediaWriter::open(const char* filename, const MFormat& fmt)
{
	if (!filename)
	{
		return false;
	}

	m_filename = filename;

	MFormatUtil::copyTo(&fmt, m_originFormat);

	if (m_originFormat.m_codec != av::MEDIA_CODEC_PS)
	{
		m_format = m_originFormat;
		openParser(m_format.m_codec);

		return m_writer->open(filename, fmt);
	}

	return true;
}

void NormMediaWriter::close()
{
	m_writer->close();

	m_startVideoPts = -1;
	m_startAudioPts = -1;
}

bool NormMediaWriter::isOpen()
{
	return m_writer->isOpen();
}

bool NormMediaWriter::getFormat(MFormat& fmt)
{
	return m_writer->getFormat(fmt);
}

bool NormMediaWriter::writeVideo(uint8_t* data, int size, int64_t pts, int flags)
{
	MPacket pkt = MPacket();
	pkt.type = MTYPE_VIDEO;
	pkt.data = data;
	pkt.size = size;
	pkt.pts = pts;
	pkt.flags = pkt.flags;
	 
	if (m_originFormat.m_codec == av::MEDIA_CODEC_PS)
	{
		return writePs(pkt);
	}

	if (m_startVideoPts < 0)
	{
		m_startVideoPts = pts;
	}

	int64_t ts = pts - m_startVideoPts;
	pkt.pts = ts;

	if (isParserOpen())
	{
		return parseWrite(pkt);
	}

	return writePacket(pkt);
}

bool NormMediaWriter::writeAudio(uint8_t* data, int size, int64_t pts, int flags)
{
	if (m_startAudioPts < 0)
	{
		m_startAudioPts = pts;
	}

	int64_t ts = pts - m_startAudioPts;

	return m_writer->writeAudio(data, size, ts, flags);
}

int NormMediaWriter::getDuration()
{
	return m_writer->getDuration();
}

int64_t NormMediaWriter::getByteLength()
{
	return m_writer->getByteLength();
}

bool NormMediaWriter::setFile(const char* filepath)
{
	return m_writer->setFile(filepath);
}

bool NormMediaWriter::setOutputAudioCodec(int codecId)
{
	return m_writer->setOutputAudioCodec(codecId);
}

void NormMediaWriter::setListener(MediaWriterListener* listener)
{
	m_writer->setListener(listener);
}

bool NormMediaWriter::setMeta(const char* key, const char* value, int size)
{
	return m_writer->setMeta(key, value, size);
}


bool NormMediaWriter::writePs(MPacket& pkt)
{
	PsParser::Packet psPacket;
	if (!m_psParser.parse(pkt.data, pkt.size, psPacket))
	{
		return false;
	}

	if (!m_format.isValid())
	{
		MediaFormat fmt = m_psParser.getFormat();
		if (fmt.isValid())
		{
			comn::AutoCritSec lock(m_cs);
			m_format = fmt;

			MFormat format = MFormat();
			MFormatUtil::copyFrom(&format, m_format);
			m_writer->open(m_filename.c_str(), format);
		}
		else
		{
			return false;
		}
	}

	MPacket outPkt = MPacket();
	outPkt.data = psPacket.data;
	outPkt.size = psPacket.length;
	outPkt.flags = psPacket.flags;
	outPkt.type = psPacket.type;

	//if (psPacket.type == av::MEDIA_TYPE_VIDEO)
	{
		outPkt.pts = psPacket.ts * 1000 / 90;
	}

	return writePacket(outPkt);
}

bool NormMediaWriter::writePacket(MPacket& pkt)
{
	return m_writer->writeVideo(pkt.data, pkt.size, pkt.pts, pkt.flags);
}


bool NormMediaWriter::openParser(int codec)
{
	AVCodecID codecID = CodecIDConvert::toAVCodecID(codec);
	m_parser = av_parser_init(codecID);
	if (m_parser == NULL)
	{
		return false;
	}

	AVCodec* pCodec = avcodec_find_decoder(AV_CODEC_ID_H264);
	m_pContext = avcodec_alloc_context3(pCodec);

	int ret = avcodec_open2(m_pContext, pCodec, NULL);
	if (ret < 0)
	{
		avcodec_close(m_pContext);
		av_free(m_pContext);
		m_pContext = NULL;
		pCodec = NULL;

		av_parser_close(m_parser);
		m_parser = NULL;

		return false;
	}

	return isParserOpen();
}

void NormMediaWriter::closeParser()
{
	if (m_pContext != NULL)
	{
		avcodec_close(m_pContext);
		av_free(m_pContext);
		m_pContext = NULL;
	}

	if (m_parser != NULL)
	{
		av_parser_close(m_parser);
		m_parser = NULL;
	}
}

bool NormMediaWriter::isParserOpen()
{
	return (m_parser != NULL);
}

bool NormMediaWriter::parseWrite(MPacket& pkt)
{
	if (isParserOpen())
	{
		uint8_t* buffer = pkt.data;
		int length = pkt.size;

		int pos = 0;
		int64_t vpts = pkt.pts;
		int64_t vdts = pkt.pts;

		do
		{
			uint8_t *pout = NULL;
			int poutLen = 0;
			int frameLen = av_parser_parse2(m_parser, m_pContext,
				&pout,
				&poutLen,
				buffer + pos,
				length - pos,
				vpts,
				vdts,
				pkt.pts);

			pos += frameLen;

			if (poutLen > 0)
			{
				MPacket packet = MPacket();
				packet.data = pout;
				packet.size = poutLen;
				packet.pts = vpts;
				packet.type = pkt.type;
				packet.flags |= pkt.flags;

				fixPacketFlags(packet);

				bool hasStartCode = H26xNaluParser::hasStartCode(packet.data, packet.size);
				assert(hasStartCode == true);

				//comn::FileUtil::write(packet.data, packet.size, "out.h264", true);
				writePacket(packet);
			}

		} while (pos < length);

		return true;
	}
	else
	{
		return writePacket(pkt);
	}
}

void NormMediaWriter::fixPacketFlags(MPacket& packet)
{
	bool isKeyFrame = false;
	av::H26xNalu nalu;
	if (m_format.m_codec == av::MEDIA_CODEC_H264)
	{
		H26xNaluParser::parseH264Nalu(packet.data, packet.size, nalu);
		isKeyFrame = (nalu.type == NALU_IFRAME || nalu.type == NALU_SPS);
	}
	else if (m_format.m_codec == av::MEDIA_CODEC_HEVC)
	{
		H26xNaluParser::parseH265Nalu(packet.data, packet.size, nalu);
		isKeyFrame = (nalu.type == NAL_IDR_N_LP || nalu.type == NAL_IDR_W_RADL || nalu.type == NAL_UNIT_VPS ||
			nalu.type == NAL_UNIT_SPS || nalu.type == NAL_UNIT_PPS);
	}

	if (isKeyFrame)
	{
		packet.flags |= av::MEDIA_FLAG_KEY;
	}
}



} /* namespace av */
